package com.github.tommas.admintpl.security;

import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.exceptions.JWTVerificationException;
import com.auth0.jwt.interfaces.DecodedJWT;
import com.github.tommas.admintpl.util.JWTUtils;
import org.apache.shiro.authc.AuthenticationInfo;
import org.apache.shiro.mgt.SecurityManager;
import org.apache.shiro.session.Session;
import org.apache.shiro.subject.PrincipalCollection;
import org.apache.shiro.subject.SimplePrincipalCollection;
import org.apache.shiro.subject.Subject;
import org.apache.shiro.subject.SubjectContext;
import org.apache.shiro.web.mgt.DefaultWebSubjectFactory;
import org.apache.shiro.web.subject.WebSubjectContext;
import org.apache.shiro.web.subject.support.WebDelegatingSubject;
import org.apache.shiro.web.util.WebUtils;
import org.springframework.util.StringUtils;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;

public class JWTAwareSubjectFactory extends DefaultWebSubjectFactory {
    public static final String TOKEN_EXCEPTION_ATTR = JWTAwareSubjectFactory.class.getName() + ".token_exception";
    private static final String AUTHORIZATION_HEADER = "Authorization";

    private Algorithm algorithm;

    public JWTAwareSubjectFactory(Algorithm algorithm) {
        super();
        this.algorithm = algorithm;
    }
    
    @Override
    public Subject createSubject(SubjectContext context) {
        WebSubjectContext wsc = (WebSubjectContext) context;
        SecurityManager securityManager = wsc.resolveSecurityManager();
        Session session = wsc.resolveSession();
        boolean sessionEnabled = wsc.isSessionCreationEnabled();
        String host = wsc.resolveHost();
        ServletRequest request = wsc.resolveServletRequest();
        ServletResponse response = wsc.resolveServletResponse();

        PrincipalCollection principals = null;
        AuthenticationInfo authcInfo = wsc.getAuthenticationInfo();
        if (authcInfo != null) {
            // try to get principals from authentication info first
            principals = authcInfo.getPrincipals();
        }

        if (principals == null) {
            // try to resole principals from jwt
            principals = getPrincipalsFromRequest(WebUtils.toHttp(request));
        }

        boolean authenticated = principals != null;

        return new WebDelegatingSubject(principals, authenticated, host, session, sessionEnabled,
                request, response, securityManager);
    }

    private PrincipalCollection getPrincipalsFromRequest(HttpServletRequest request) {
        String authorizationHeader = request.getHeader(AUTHORIZATION_HEADER);
        if (!StringUtils.hasLength(authorizationHeader)) {
            return null;
        }

        String[] authTokens = authorizationHeader.split(" ");
        if (authTokens.length < 2) {
            return null;
        }

        String token = authTokens[1];
        try {
            DecodedJWT decoded = JWTUtils.decodeToken(token, algorithm);

            String username = decoded.getClaim(JWTUtils.USERNAME_CLAIM).asString();
            Integer userId = decoded.getClaim(JWTUtils.USER_ID_CLAIM).asInt();
            if (!StringUtils.hasLength(username) || userId == null) {
                throw new RuntimeException(String.format("Invalid token: %s", token));
            }

            SimplePrincipalCollection pc = new SimplePrincipalCollection();
            pc.add(username, AccountRealm.REALM_NAME);
            pc.add(userId, AccountRealm.REALM_NAME);

            return pc;
        } catch (JWTVerificationException e) {
            request.setAttribute(TOKEN_EXCEPTION_ATTR, e);
        }

        return null;
    }
}